
package executestage

import chisel3.util._
import chisel3._
import executestage.ALUOpCode._
object ALUOpCode {
  val ALU_Width = 17.U.getWidth
  val ALU_NON = 0.U(ALU_Width.W)
  val ALU_ADD = 1.U(ALU_Width.W)
  val ALU_ADDU = 2.U(ALU_Width.W)
  val ALU_SUB = 3.U(ALU_Width.W)
  val ALU_SUBU = 4.U(ALU_Width.W)
  val ALU_OR = 5.U(ALU_Width.W)
  val ALU_XOR = 6.U(ALU_Width.W)
  val ALU_AND = 7.U(ALU_Width.W)
  val ALU_NOR = 8.U(ALU_Width.W)
  val ALU_RS = 9.U(ALU_Width.W)
  val ALU_RT = 10.U(ALU_Width.W)
  val ALU_IMM = 11.U(ALU_Width.W)
  val ALU_SLL = 12.U(ALU_Width.W)
  val ALU_SRL = 13.U(ALU_Width.W)
  val ALU_SRA = 14.U(ALU_Width.W)
  val ALU_SLT = 15.U(ALU_Width.W)
  val ALU_SLTU = 16.U(ALU_Width.W)
  val ALU_Link = 17.U(ALU_Width.W)
}
class ALUInBus extends Bundle {
  val rsData = UInt(32.W)
  val rtData = UInt(32.W)
  val PC = UInt(32.W)
  val imm32 = UInt(32.W)
  val opCode = UInt(ALU_Width.W)
  val isReg = Bool()
}
class ALU extends Module {
  val io = IO(new Bundle {
    val in = Input(new ALUInBus)
    val out = Output(UInt(32.W))
  })
  val in = io.in
  // add and sub
  val AS_A = in.rsData.asSInt()
  val AS_B = Mux(in.isReg, in.rtData, in.imm32).asSInt()
  val addAnsTmp = AS_A +& AS_B
  val subAnsTmp = AS_A -& AS_B
  assert(addAnsTmp.getWidth == 33)
  assert(subAnsTmp.getWidth == 33)
  val addOverFlow = addAnsTmp(32) =/= addAnsTmp(31)
  val subOverFlow = subAnsTmp(32) =/= subAnsTmp(31)
  val addAns = addAnsTmp(31, 0).asUInt()
  val subAns = subAnsTmp(31, 0).asUInt()
  //Bit Cal
  val BC_A = in.rsData
  val BC_B = Mux(in.isReg, in.rtData, in.imm32)
  val orAns = BC_A | BC_B
  val xorAns = BC_A ^ BC_B
  val andAns = BC_A & BC_B
  val norAns = (~orAns).asUInt()
  //Shift Family
  val ST_A = Mux(in.isReg, in.rtData, in.imm32)(4, 0)
  val ST_B = in.rtData
  val sllAns = (ST_B << ST_A) (31, 0)
  val srlAns = (ST_B >> ST_A).asUInt()
  val sraAns = (ST_B.asSInt() >> ST_A).asUInt()
  assert(sllAns.getWidth == 32 && srlAns.getWidth == 32 && sraAns.getWidth == 32)
  assert(ST_A.getWidth == 5)
  //CMP Family
  val CMP_A = in.rsData
  val CMP_B = Mux(in.isReg, in.rtData, in.imm32)
  val sltAns = Cat(0.U(31.W), CMP_A.asSInt() < CMP_B.asSInt())
  val sltuAns = Cat(0.U(31.W), CMP_A.asUInt() < CMP_B.asUInt())

  val linkAns = in.PC + 8.U

  val opCode = in.opCode
  io.out := MuxCase(0.U, Seq(
    (opCode === ALU_ADD || opCode === ALU_ADDU) -> addAns,
    (opCode === ALU_SUB || opCode === ALU_SUBU) -> subAns,
    (opCode === ALU_OR) -> orAns,
    (opCode === ALU_XOR) -> xorAns,
    (opCode === ALU_AND) -> andAns,
    (opCode === ALU_NOR) -> norAns,
    (opCode === ALU_RS) -> in.rsData,
    (opCode === ALU_RT) -> in.rtData,
    (opCode === ALU_IMM) -> in.imm32,
    (opCode === ALU_SLL) -> sllAns,
    (opCode === ALU_SRL) -> srlAns,
    (opCode === ALU_SRA) -> sraAns,
    (opCode === ALU_SLT) -> sltAns,
    (opCode === ALU_SLTU) -> sltuAns,
    (opCode === ALU_Link) -> linkAns
  ))
}
